from pydantic import BaseModel, Field
from langchain.prompts import PromptTemplate
from langchain_deepseek import ChatDeepSeek

# 定义结构化输出模型
class Date(BaseModel):
    year: int = Field(description="Year")
    month: int = Field(description="Month")
    day: int = Field(description="Day")
    era: str = Field(description="BC or AD")

# 初始化 LLM
llm = ChatDeepSeek(model="deepseek-chat", temperature=0)

# 定义结构化输出模型
structured_llm = llm.with_structured_output(Date)

template = """提取用户输入中的日期。
用户输入:
{query}"""

prompt = PromptTemplate(
    template = template
)

query = "2024年十二月23日天气晴..."

input_prompt = prompt.format_prompt(query = query)
result = structured_llm.invoke(input_prompt)
print(type(result))
print(result)


